#!/usr/bin/env python
from astropy.io import fits
import os
from argparse import ArgumentParser
from astropy.time import Time
from glob import glob

def get_header(hdr,obstype,chitype,chinum,systemp,time,idx):
    hdr0 = hdr.copy()
    isotime = time.iso
    hdr0['DATE'] = isotime
    hdr0['DATE-OBS'] = isotime
    hdr0['DATETAG'] = time2tag(isotime)
    hdr0['OBSTYPE'] = obstype
    hdr0['CHITYPE'] = chitype
    hdr0['CHINUM'] = chinum
    hdr0['EXPSTART'] = time.mjd
    if obstype == 'bias':
        hdr0['EXPTIME'] = 0.03
        hdr0['DARKTIME'] = 2
    elif obstype == 'dark':
        hdr0['EXPTIME'] = 1000
        hdr0['DARKTIME'] = 1000
    hdr0['WAVELEN'] = 0.0
    hdr0['IMGINDEX'] = 0
    hdr0['SEQNUM'] = idx
    hdr0['SYSTEMP'] = float(systemp)
    hdr0['BUNIT'] = 'adu'
    hdr0['GLASSTYP'] = obstype+"-"+str(idx)
    return hdr0

def tag2time(tag):
    time = tag[0:4]+"-"+tag[4:6]+"-"+tag[6:8]+" "+tag[9:11]+":"+tag[11:13]+":"+tag[13:15]
    return time
def time2tag(time):
    tag = time[0:4]+time[5:7]+time[8:10]+"T"+time[11:13]+time[14:16]+time[17:19]
    return tag
    
def add_header(in_dir,out_dir,bias_b_hdr,dark_hdr,bias_a_hdr,tag,chitype,chinum,systemp):
    files = glob(in_dir+"/*.fits")
    tag0 = tag
    time0 = tag2time(tag0)
    print("time0",time0)
    bias_b_time = Time(time0)
    bias_a_time = Time(Time(time0).mjd+3/24,format='mjd')
    taga = time2tag(bias_a_time.iso)
    dark_time = Time(Time(time0).mjd+1/24,format='mjd')
    tagd = time2tag(dark_time.iso)
    bias_b_idx = 1
    bias_a_idx = 1
    dark_idx = 1
    if not os.path.exists(out_dir+"/dark_"+tag0[0:8]):
        os.system("mkdir "+out_dir+"/dark_"+tag0[0:8])
    if not os.path.exists(out_dir+"/dark_"+tag0[0:8]+"/"+systemp+"K"):
        os.system("mkdir "+out_dir+"/dark_"+tag0[0:8]+"/"+systemp+"K")
    dir_hdr =  out_dir+"/dark_"+tag0[0:8]+"/"+systemp+"K"
    if not os.path.exists(dir_hdr+"/bias_"+tag0):
        os.system("mkdir "+dir_hdr+"/bias_"+tag0)
    if not os.path.exists(dir_hdr+"/bias_"+taga):
        os.system("mkdir "+dir_hdr+"/bias_"+taga)
    if not os.path.exists(dir_hdr+"/dark_"+tagd):
        os.system("mkdir "+dir_hdr+"/dark_"+tagd)
    for i in range(len(files)):
        f = files[i]
        data = fits.getdata(f)
        hdr = fits.open(f)[0].header
        fname = f.split('/')[-1]
        if fname[:len(bias_b_hdr)]==bias_b_hdr:
            hdr = get_header(hdr,'bias',chitype,chinum,systemp,bias_b_time,bias_b_idx)
            fits.writeto(dir_hdr+"/bias_"+tag0+"/bias-"+str(bias_b_idx)+".fits",data=data,header=hdr,overwrite=True)
            bias_b_idx += 1
            bias_b_time = Time(bias_b_time.mjd+1/60/24,format='mjd')
        elif fname[:len(bias_a_hdr)]==bias_a_hdr:
            hdr = get_header(hdr,'bias',chitype,chinum,systemp,bias_a_time,bias_a_idx)
            fits.writeto(dir_hdr+"/bias_"+taga+"/bias-"+str(bias_a_idx)+".fits",data=data,header=hdr,overwrite=True)
            bias_a_idx += 1
            bias_a_time = Time(bias_a_time.mjd+1/60/24,format='mjd')      
        elif fname[:len(dark_hdr)]==dark_hdr:
            hdr = get_header(hdr,'dark',chitype,chinum,systemp,dark_time,dark_idx)
            fits.writeto(dir_hdr+"/dark_"+tagd+"/dark-"+str(dark_idx)+".fits",data=data,header=hdr,overwrite=True)
            dark_idx += 1
            dark_time = Time(dark_time.mjd+1/60/24,format='mjd')
    return 
if __name__ == '__main__':
    parser = ArgumentParser()
    parser.add_argument('indir',type=str) #178K_fits
    parser.add_argument('outdir', type=str) #data
    parser.add_argument('--bias_b_hdr','-bh',type=str)
    parser.add_argument('--dark_hdr','-dh',type=str)
    parser.add_argument('--bias_a_hdr','-ah',type=str)
    parser.add_argument('--tag','-t',type=str)
    parser.add_argument('--chitype','-ct',type=str)
    parser.add_argument('--chinum','-ci',type=str)
    parser.add_argument('--systemp','-st',type=str)
    args = parser.parse_args()
    add_header(args.indir,args.outdir,args.bias_b_hdr,args.dark_hdr,args.bias_a_hdr,args.tag,args.chitype
    ,args.chinum,args.systemp)
